package com.mcp.mcpserver.service;

import com.mcp.mcpserver.entity.DatabaseInstance;
import com.mcp.mcpserver.model.McpResponse;
import com.zaxxer.hikari.HikariConfig;
import com.zaxxer.hikari.HikariDataSource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.jdbc.core.namedparam.NamedParameterJdbcTemplate;
import org.springframework.stereotype.Service;

import javax.sql.DataSource;
import java.sql.SQLException;
import java.util.List;
import java.util.Map;
import java.util.Set;

@Service
public class DynamicQueryService {
    private static final Logger logger = LoggerFactory.getLogger(DynamicQueryService.class);

    // 定义危险操作关键字
    private static final Set<String> DANGEROUS_OPERATIONS = Set.of(
            "DROP", "TRUNCATE", "DELETE", "ALTER", "CREATE", "SHUTDOWN"
    );

    public McpResponse executeQuery(DatabaseInstance instance, String sql, Map<String, Object> params) {
        // 创建动态数据源
        DataSource dataSource = createDynamicDataSource(instance);

        try {
            // 验证数据源连接
            testConnection(dataSource);

            // 执行查询
            return executeQueryWithDataSource(dataSource, sql, params);
        } catch (SQLException e) {
            logger.error("Database connection test failed", e);
            return McpResponse.error("CONNECTION_ERROR", "Failed to connect to database: " + e.getMessage());
        } finally {
            // 关闭临时数据源
            if (dataSource instanceof HikariDataSource) {
                ((HikariDataSource) dataSource).close();
            }
        }
    }

    private DataSource createDynamicDataSource(DatabaseInstance instance) {
        HikariConfig config = new HikariConfig();

        // 根据数据库类型构建JDBC URL
        String jdbcUrl = buildJdbcUrl(instance);
        logger.debug("Received request for dbRequestBody : {}",
                jdbcUrl);
        config.setJdbcUrl(jdbcUrl);
        config.setUsername(instance.getDbUsername());
        config.setPassword(instance.getDbPassword());
        config.setDriverClassName(getDriverClassName(instance.getDbType()));

        // 连接池配置
        config.setMaximumPoolSize(5);
        config.setMinimumIdle(1);
        config.setConnectionTimeout(30000);
        config.setIdleTimeout(600000);
        config.setMaxLifetime(1800000);

        return new HikariDataSource(config);
    }

    private String buildJdbcUrl(DatabaseInstance instance) {
        switch (instance.getDbType().toLowerCase()) {
            case "mysql":
                return String.format("jdbc:mysql://%s:%d/%s?useUnicode=true&characterEncoding=utf8&useSSL=false&serverTimezone=UTC&allowPublicKeyRetrieval=true",
                        instance.getDbHost(), instance.getDbPort(), instance.getDbDatabase());
            case "oceanbase":
                StringBuilder obUrl = new StringBuilder();
                obUrl.append("jdbc:oceanbase://")
                        .append(instance.getDbHost())
                        .append(":")
                        .append(instance.getDbPort())
                        .append("/")
                        .append(instance.getDbDatabase());
                String tenant = instance.getDbTenant();
                // TODO 添加租户参数
                if (tenant != null && !tenant.trim().isEmpty()) {
                    obUrl.append("?tenant=").append(tenant);
                }
                // 添加性能参数
                boolean hasQuery = obUrl.indexOf("?") > 0;
                obUrl.append(hasQuery ? "&" : "?")
                        .append("rewriteBatchedStatements=true")
                        .append("&allowMultiQueries=true")
                        .append("&useLocalSessionState=true")
                        .append("&characterEncoding=utf-8");

                return obUrl.toString();
            case "postgresql":
                return String.format(
                        "jdbc:postgresql://%s:%d/%s",
                        instance.getDbHost(), instance.getDbPort(), instance.getDbDatabase()
                );

            default:
                throw new IllegalArgumentException("Unsupported database type: " + instance.getDbType());
        }
    }

    private String getDriverClassName(String dbType) {
        switch (dbType.toLowerCase()) {
            case "mysql":
                return "com.mysql.cj.jdbc.Driver";
            case "oceanbase":
                return "com.oceanbase.jdbc.Driver";
            case "postgresql":
                return "org.postgresql.Driver";
            default:
                throw new IllegalArgumentException("Unsupported database type: " + dbType);
        }
    }

    private void testConnection(DataSource dataSource) throws SQLException {
        try (var connection = dataSource.getConnection()) {
            if (!connection.isValid(5)) {
                throw new SQLException("Connection test failed");
            }
        }
    }


    // 添加 SQL 注入防护方法
    private void validateSql(String sql) {
        if (sql == null || sql.trim().isEmpty()) {
            throw new IllegalArgumentException("SQL query cannot be empty");
        }

        // 检查危险操作
        String upperSql = sql.toUpperCase();
        // 检查危险操作
        for (String operation : DANGEROUS_OPERATIONS) {
            if (upperSql.contains(operation) &&
                    !upperSql.contains("/* SAFE_OPERATION */")) {
                throw new SecurityException("Potentially dangerous SQL operation detected: " + operation);
            }
        }

        // 检查注入特征
        if (upperSql.contains(";") ||
                upperSql.contains("--") ||
                upperSql.contains("/*")) {
            throw new SecurityException("Query contains suspicious characters");
        }
    }



    private McpResponse executeQueryWithDataSource(DataSource dataSource, String sql, Map<String, Object> params) {
        // 先进行 SQL 安全验证
        // validateSql(sql);
        NamedParameterJdbcTemplate jdbcTemplate = new NamedParameterJdbcTemplate(dataSource);

        try {
            String upperSql = sql.trim().toUpperCase();

            // 处理元数据查询 (SHOW, DESCRIBE 等)
            if (upperSql.startsWith("SHOW") ||
                    upperSql.startsWith("DESCRIBE") ||
                    upperSql.startsWith("DESC") ||
                    upperSql.startsWith("EXPLAIN")) {

                // 使用简单查询（不需要参数化）
                List<Map<String, Object>> results = jdbcTemplate.getJdbcTemplate().queryForList(sql);
                return McpResponse.success(results);
            }
            // 处理 SELECT 查询
            else if (sql.trim().toUpperCase().startsWith("SELECT")) {
                // 查询操作
                List<Map<String, Object>> results = jdbcTemplate.queryForList(sql, params);
                return McpResponse.success(results);
            } else {
                // 更新操作（INSERT/UPDATE/DELETE）
                int affectedRows = jdbcTemplate.update(sql, params);
                return McpResponse.success(Map.of("affected_rows", affectedRows));
            }
        } catch (Exception e) {
            logger.error("Error executing SQL query", e);
            return McpResponse.error("QUERY_ERROR", "Error executing SQL: " + e.getMessage());
        }
    }
}